import os
import glob
import random
import re
from tqdm import tqdm
import json

def process_java_files():
    """
    处理dataset文件夹中的Java文件，移除注释行并创建代码片段列表
    每个列表包含3个代码片段，每个片段包含1-6行随机代码
    """
    # 获取dataset目录下所有Java文件
    java_files = glob.glob('dataset/**/*.java', recursive=True)
    
    # 大列表，存储所有处理结果
    all_code_triplets = []
    
    # 使用tqdm展示遍历进度
    for file_path in tqdm(java_files, desc="处理Java文件"):
        file_triplets = process_single_file(file_path)
        all_code_triplets.extend(file_triplets)
    
    return all_code_triplets

def process_single_file(file_path):
    """处理单个Java文件并返回代码片段三元组列表"""
    with open(file_path, 'r', encoding='utf-8') as file:
        # 读取文件内容
        lines = file.readlines()
        
        # 删除注释行，使用tqdm展示每行处理进度
        clean_lines = []
        for line in tqdm(lines, desc=f"处理文件 {os.path.basename(file_path)}", leave=False):
            # 检查是否是注释行 (以"//"开头或连续空格后有"//")
            if not re.match(r'^\s*\/\/', line):
                clean_lines.append(line)
        
        # 如果处理后的文件为空，返回空列表
        if not clean_lines:
            return []
        
        # 从文件中创建代码三元组
        file_triplets = []
        
        # 确定能创建多少个完整的三元组
        max_triplets = len(clean_lines) // 3  # 只是一个估计值，因为片段长度是随机的
        
        for _ in tqdm(range(max_triplets), desc=f"生成 {os.path.basename(file_path)} 的代码三元组", leave=False):
            # 尝试创建一个包含3个代码片段的列表
            triplet = []
            
            # 创建3个代码片段
            for _ in range(3):
                if not clean_lines:  # 如果没有更多行可用
                    break
                
                # 随机决定片段长度 (1-6行)
                segment_length = min(random.randint(1, 6), len(clean_lines))
                
                # 随机选择起始位置
                if len(clean_lines) > segment_length:
                    start_idx = random.randint(0, len(clean_lines) - segment_length)
                else:
                    start_idx = 0
                
                # 提取代码片段
                segment = ''.join(clean_lines[start_idx:start_idx + segment_length])
                triplet.append(segment)
                
                # 移除已使用的行(可选，取决于是否允许重叠)
                # 如果不希望同一个文件中的代码片段有重叠，取消下面这行的注释
                # clean_lines = clean_lines[:start_idx] + clean_lines[start_idx + segment_length:]
            
            # 只有当成功创建了3个片段时，才添加到结果中
            if len(triplet) == 3:
                file_triplets.append(triplet)
            else:
                break
        
        return file_triplets

def create_completion_dataset(triplets):
    """
    将代码三元组转换为挖空补全任务的字典列表
    每个字典包含instruction和output字段
    随机选择A、B或C进行挖空
    """
    completion_data = []
    
    for triplet in tqdm(triplets, desc="生成补全任务"):
        # 确保三元组长度为3
        if len(triplet) != 3:
            continue
            
        # 随机选择要挖空的片段索引 (0, 1, 或 2，对应A、B、C)
        blank_idx = random.randint(0, 2)
        
        # 保存被挖空的片段作为答案，并对大括号进行转义
        answer = triplet[blank_idx].strip()
        
        # 创建带有挖空的代码
        code_pieces = []
        for i, segment in enumerate(triplet):
            if i == blank_idx:
                code_pieces.append("< | fim_hole | >")
            else:
                code_pieces.append(segment.strip())
        
        # 将挖空后的代码片段连接起来
        code_string = "".join(code_pieces)
        
        # 创建指令和输出字典
        completion_item = {
            "input": f"< | fim_begin | >{code_string}< | fim_end | >",
            "output": f"{answer};",
        }
        
        completion_data.append(completion_item)
    
    return completion_data

class SafeJSONEncoder(json.JSONEncoder):
    def encode(self, obj):
        """确保所有字符串中的特殊JSON字符都被正确转义"""
        return super().encode(obj)
    
    def iterencode(self, obj, _one_shot=False):
        """自定义迭代编码确保转义特殊字符"""
        return super().iterencode(obj, _one_shot)

def save_dataset(data, filename="completion_dataset.json"):
    """使用安全的JSON编码器保存数据集"""
    with open(filename, 'w', encoding='utf-8') as f:
        json.dump(data, f, ensure_ascii=False, indent=2, cls=SafeJSONEncoder)
    print(f"数据集已保存至 {filename}，共 {len(data)} 个样本")

def main():
    # 处理Java文件获取代码三元组
    code_triplets = process_java_files()
    print(f"总共生成了 {len(code_triplets)} 个代码片段三元组")
    
    # 创建代码补全任务数据集
    completion_dataset = create_completion_dataset(code_triplets)
    
    # 输出示例
    if completion_dataset:
        print("\n生成的数据集示例:")
        for i in range(min(3, len(completion_dataset))):
            print(f"\n示例 {i+1}:")
            print(json.dumps(completion_dataset[i], ensure_ascii=False, indent=2))
    
    # 保存数据集
    save_dataset(completion_dataset)
    
    return completion_dataset

if __name__ == "__main__":
    main()